package beautiful.butterfly.drds.data_exchange.job;


import beautiful.butterfly.drds.data_exchange.constant.Constants;
import beautiful.butterfly.drds.data_exchange.util.Configuration;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.Validate;

import java.util.*;


public final class TaskAssigns
{
    private TaskAssigns()
    {
    }

    /**
     * 公平的分配 task 到对应的 taskGroup 中。
     * 公平体现在：会考虑 task 中对资源负载作的 load 标识进行更均衡的作业分配操作。
     * TODO 具体文档举例说明
     */
    public static List<Configuration> assignFairly(Configuration configuration, int channelNumber, int channelsPerTaskGroup)
    {
        Validate.isTrue(configuration != null, "框架获得的 job 不能为 null.");

        List<Configuration> jobConfigurationList = configuration.getListConfiguration(Constants.task_list);
        Validate.isTrue(jobConfigurationList.size() > 0, "框架获得的切分后的 job 无内容.");

        Validate.isTrue(channelNumber > 0 && channelsPerTaskGroup > 0,
                "每个channel的平均task数[averTaskPerChannel]，channel数目[channelNumber]，每个taskGroup的平均channel数[channelsPerTaskGroup]都应该为正数");

        int taskGroupNumber = (int) Math.ceil(1.0 * channelNumber / channelsPerTaskGroup);

        Configuration configuration0 = jobConfigurationList.get(0);

        String readerResourceMark = configuration0.getString(Constants.job_reader_parameter + "." + Constants.load_balance_resource_mark);
        String writerResourceMark = configuration0.getString(Constants.job_writer_parameter + "." + Constants.load_balance_resource_mark);

        boolean hasLoadBalanceResourceMark = StringUtils.isNotBlank(readerResourceMark) || StringUtils.isNotBlank(writerResourceMark);

        if (!hasLoadBalanceResourceMark)
        {
            // fake 一个固定的 key 作为资源标识（在 reade 或者 write 上均可，此处选择在 reade 上进行 fake）
            for (Configuration configuration1 : jobConfigurationList)
            {
                configuration1.set(Constants.job_reader_parameter + "." +
                        Constants.load_balance_resource_mark, "aFakeResourceMarkForLoadBalance");
            }
            // 是为了避免某些插件没有设置 资源标识 而进行了一次随机打乱操作
            Collections.shuffle(jobConfigurationList, new Random(System.currentTimeMillis()));
        }

        LinkedHashMap<String, List<Integer>> shuffleResourceMarkToTaskIdListMap = getShuffleResourceMarkToTaskIdListMap(jobConfigurationList);
        List<Configuration> taskGroupConfigurationList = assign(shuffleResourceMarkToTaskIdListMap, configuration, taskGroupNumber);

        // 调整 每个 taskGroup 对应的 RecordQueue 个数（属于优化范畴）
        adjustChannelNumPerTaskGroup(taskGroupConfigurationList, channelNumber);
        return taskGroupConfigurationList;
    }


    /**
     * 根据task 配置，获取到：
     * 资源名称 --> taskId(List) 的 map 映射关系
     */
    private static LinkedHashMap<String, List<Integer>> getShuffleResourceMarkToTaskIdListMap(List<Configuration> configurationList)
    {
        // key: resourceMark, value: taskId
        LinkedHashMap<String, List<Integer>> readerResourceMarkToTaskIdListMap = new LinkedHashMap<String, List<Integer>>();
        LinkedHashMap<String, List<Integer>> writerResourceMarkToTaskIdListMap = new LinkedHashMap<String, List<Integer>>();

        for (Configuration configuration : configurationList)
        {
            int taskId = configuration.getInt(Constants.task_id);
            // 把 readerResourceMark 加到 readerResourceMarkAndTaskIdMap 中
            String readerResourceMark = configuration.getString(Constants.job_reader_parameter + "." + Constants.load_balance_resource_mark);
            if (readerResourceMarkToTaskIdListMap.get(readerResourceMark) == null)
            {
                readerResourceMarkToTaskIdListMap.put(readerResourceMark, new LinkedList<Integer>());
            }
            readerResourceMarkToTaskIdListMap.get(readerResourceMark).add(taskId);

            // 把 writerResourceMark 加到 writerResourceMarkAndTaskIdMap 中
            String writerResourceMark = configuration.getString(Constants.job_writer_parameter + "." + Constants.load_balance_resource_mark);
            if (writerResourceMarkToTaskIdListMap.get(writerResourceMark) == null)
            {
                writerResourceMarkToTaskIdListMap.put(writerResourceMark, new LinkedList<Integer>());
            }
            writerResourceMarkToTaskIdListMap.get(writerResourceMark).add(taskId);
        }

        if (readerResourceMarkToTaskIdListMap.size() >= writerResourceMarkToTaskIdListMap.size())
        {
            // 采用 reade 对资源做的标记进行 shuffle
            return readerResourceMarkToTaskIdListMap;
        } else
        {
            // 采用 write 对资源做的标记进行 shuffle
            return writerResourceMarkToTaskIdListMap;
        }
    }


    /**
     * /**
     * 需要实现的效果通过例子来说是：
     * <pre>
     * a 库上有表：0, 1, 2
     * a 库上有表：3, 4
     * c 库上有表：5, 6, 7
     *
     * 如果有 4个 taskGroup
     * 则 assign 后的结果为：
     * taskGroup-0: 0,  4,
     * taskGroup-1: 3,  6,
     * taskGroup-2: 5,  2,
     * taskGroup-3: 1,  7
     *
     * </pre>
     */
    private static List<Configuration> assign(LinkedHashMap<String, List<Integer>> resourceMarkToTaskIdListMap, Configuration jobConfiguration, int taskGroupNumber)
    {
        List<Configuration> jobConfigurationList = jobConfiguration.getListConfiguration(Constants.task_list);
        Configuration taskGroupTemplate = jobConfiguration.clone();
        taskGroupTemplate.remove(Constants.task_list);

        List<Configuration> configurationList = new LinkedList<Configuration>();

        List<List<Configuration>> taskGroupConfigurationListList = new ArrayList<List<Configuration>>(taskGroupNumber);
        for (int i = 0; i < taskGroupNumber; i++)
        {
            taskGroupConfigurationListList.add(new LinkedList<Configuration>());
        }

        int maxLength = -1;

        List<String> resourceMarkList = new ArrayList<String>();
        for (Map.Entry<String, List<Integer>> entry : resourceMarkToTaskIdListMap.entrySet())
        {
            resourceMarkList.add(entry.getKey());
            if (entry.getValue().size() > maxLength)
            {
                maxLength = entry.getValue().size();
            }
        }

        int taskGroupIndex = 0;
        for (int i = 0; i < maxLength; i++)
        {
            for (String resourceMark : resourceMarkList)
            {
                if (resourceMarkToTaskIdListMap.get(resourceMark).size() > 0)
                {
                    int taskId = resourceMarkToTaskIdListMap.get(resourceMark).get(0);
                    taskGroupConfigurationListList.get(taskGroupIndex % taskGroupNumber).add(jobConfigurationList.get(taskId));
                    resourceMarkToTaskIdListMap.get(resourceMark).remove(0);
                    //
                    taskGroupIndex++;


                }
            }
        }

        Configuration configuration;
        for (int i = 0; i < taskGroupNumber; i++)
        {
            configuration = taskGroupTemplate.clone();
            configuration.set(Constants.container_taskgroup_id, i);
            configuration.set(Constants.task_list, taskGroupConfigurationListList.get(i));


            configurationList.add(configuration);
        }

        return configurationList;
    }

    private static void adjustChannelNumPerTaskGroup(List<Configuration> configurationList, int channelNumber)
    {
        int taskGroupNumber = configurationList.size();
        int avgChannelsPerTaskGroup = channelNumber / taskGroupNumber;
        int remainderChannelCount = channelNumber % taskGroupNumber;
        // 表示有 remainderChannelCount 个 taskGroup,其对应 RecordQueue 个数应该为：avgChannelsPerTaskGroup + 1；
        // （taskGroupNumber - remainderChannelCount）个 taskGroup,其对应 RecordQueue 个数应该为：avgChannelsPerTaskGroup

        int i = 0;
        for (; i < remainderChannelCount; i++)
        {
            configurationList.get(i).set(Constants.container_taskgroup_channel, avgChannelsPerTaskGroup + 1);
        }

        for (int j = 0; j < taskGroupNumber - remainderChannelCount; j++)
        {
            configurationList.get(i + j).set(Constants.container_taskgroup_channel, avgChannelsPerTaskGroup);
        }
    }
}
